from langchain_ollama import ChatOllama
from typing import Optional, List
from typing_extensions import Annotated, TypedDict

# llm = ChatOllama(
#     model="deepseek-r1:14b",
#     base_url="http://localhost:11434",
#     temperature=0.3,
#     num_predict=2048,
#     top_k=50,
#     top_p=0.9,
#     stop=["<|EOT|>"],
#     num_gpu=99,
#     num_thread=8,
# )
from langchain_ollama import ChatOllama
# 1. 初始化 Ollama 模型连接
llm = ChatOllama(model="deepseek-r1:14b", validate_model_on_init=True)

class ShortVideoScript(TypedDict):
    """抖音短视频脚本结构"""
    hook: Annotated[Optional[str], "黄金开头：5秒内吸引眼球的强冲击力内容"]
    setup: Annotated[Optional[str], "背景铺垫：交代场景和问题，引发共鸣"]
    core: Annotated[List[str], "核心内容：分步骤或分要点呈现核心信息（至少3点）"]
    emotion: Annotated[Optional[str], "情绪升华：激发用户情感共鸣（如感动、好奇、愤怒）"]
    cta: Annotated[Optional[str], "行动引导：明确引导用户点赞、关注或评论"]
    easter_egg: Annotated[Optional[str], "结尾彩蛋：意外反转或隐藏福利，提升完播率"]

structured_llm = llm.with_structured_output(ShortVideoScript)

result = structured_llm.invoke(
    f"请为主题'职场塑料姐妹花'创作一个犀利幽默风格的抖音短视频脚本，包含黄金开头、背景铺垫、核心内容、情绪升华、行动引导和结尾彩蛋六个部分。"
)

print(result)